/*
 *  This file is part of AtlasMapper server and clients.
 *
 *  Copyright (C) 2012 Australian Institute of Marine Science
 *
 *  Contact: Gael Lafond <g.lafond@aims.org.au>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package au.gov.aims.colouriseimage;

import javax.imageio.ImageIO;
import java.awt.Graphics2D;
import java.awt.image.BufferedImage;
import java.awt.image.IndexColorModel;
import java.awt.image.Raster;
import java.awt.image.WritableRaster;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.logging.Level;
import java.util.logging.Logger;

public class Colouriser {
	private static final Logger LOGGER = Logger.getLogger(Colouriser.class.getName());

	public static BufferedImage getBufferedImage(URL imageUrl) throws IOException {
		return ImageIO.read(imageUrl);
	}

	public static BufferedImage getBufferedImage(InputStream imageStream) throws IOException {
		return ImageIO.read(imageStream);
	}

	/**
	 * Multiply filter. Works with PNG, JPEG but not GIF.
	 * #FF0000 keeps only red, #FF00FF keeps red and blue, etc.
	 *
	 * NOTE: I could get the YIQ colour (to transform to grey scale) before multiplying, which would ack
	 * more like expected, but this service is intended to be used with KML, which use grey scale
	 * PNG markers.
	 * @param image
	 * @param red
	 * @param green
	 * @param blue
	 * @return
	 */
	public static BufferedImage colourImage(BufferedImage image, int red, int green, int blue) {
		return Colouriser.colourImage(image, red, green, blue, "UNKNOWN");
	}

	/**
	 * This method takes an image and alter its colour.
	 * Support pretty much any type of image, except animated images (like animated GIF). Those are container and can
	 * not be modified the same way as static images.
	 * For more info about how to manipulate animated GIF:
	 *     http://elliot.kroo.net/software/java/GifSequenceWriter/
	 * @param image The original image (may be altered for better performances)
	 * @param red The red part (0 to 255) of the colour used to colourise the image.
	 * @param green The green part (0 to 255) of the colour used to colourise the image.
	 * @param blue The blue part (0 to 255) of the colour used to colourise the image.
	 * @param source Optional parameter used for debugging only.
	 * @return A colourised image.
	 */
	public static BufferedImage colourImage(BufferedImage image, int red, int green, int blue, String source) {
		int imageType = image.getType();

		if (image == null) {
			return null;
		}
		LOGGER.log(Level.FINE, "Colourising an image of type: {0}   Source: {1}", new Object[]{imageType, source});

		red = Colouriser.normaliseColour(red);
		green = Colouriser.normaliseColour(green);
		blue = Colouriser.normaliseColour(blue);
		if (red == 255 && green == 255 && blue == 255) {
			return image;
		}

		// If the image is not supported, convert it to a supported format.
		if (!isSupported(imageType)) {
			LOGGER.log(Level.FINE, "Unsupported image of type: {0}   Source: {1}. Converting to TYPE_4BYTE_ABGR (PNG with alpha)", new Object[]{imageType, source});
			image = Colouriser.convert(image, BufferedImage.TYPE_4BYTE_ABGR);
			imageType = image.getType();
		}

		int width = image.getWidth();
		int height = image.getHeight();

		try {
			if (imageType == BufferedImage.TYPE_BYTE_INDEXED) {
				// Indexed images (the pixels represent the index of its colour described in the index colour model - the colour map);
				//     GIF, Indexed PNG, etc.
				// Since the colours are described in the index colour model, the raster do not need to be modified.
				//     Only the index colour model has to be modified.
				Raster raster = image.getData();
				IndexColorModel indexedModel = (IndexColorModel) image.getColorModel();
				int size = indexedModel.getMapSize();
				byte[] reds = new byte[size], greens = new byte[size], blues = new byte[size];
				indexedModel.getReds(reds);
				indexedModel.getGreens(greens);
				indexedModel.getBlues(blues);

				// All colour using this index are displayed as transparent pixels
				int transparentPixel = indexedModel.getTransparentPixel();

				// Compute some values here instead of inside the loop, to save valuable CPU time
				float ratioRed = red / (float)0xFF,
					ratioGreen = green / (float)0xFF,
					ratioBlue = blue / (float)0xFF;

				// This will also modify the value of the transparent pixel. That doesn't matter since it is
				// referenced by its index, not its colour.
				if (reds.length == greens.length && reds.length == blues.length) {
					for (int i=0; i<reds.length; i++) {
						// In java, Bytes are always signed.
						// Mathematical operation on Bytes greater than 0x7F result in errors.
						// To work safely with Bytes, the value has to be "designed" using the binary and trick.
						// Example of the binary and trick using the Byte 0x88:
						//     byte a = (byte)0x88;                 // Unsigned 0x88 = 136
						//     byte wrong = (byte)(a / 2);          //  wrong = 0xc4 = -60
						//     byte right = (byte)((a & 0xFF) / 2); //  right = 0x44 =  68
						// Memory representation of byte 88:
						//     The real unsigned Byte 88:                                 10111000 (It is impossible to accurately represent this in Java)
						//     Java representation of Byte 88: 11111111 11111111 11111111 10111000 (Bytes are stored in 16 bits in memory)
						//     Binary representation of 0xFF:  00000000 00000000 00000000 11111111
						//     After binary and with 0xFF:     00000000 00000000 00000000 10111000
						// Math example of the problem:
						//     In C, C++ and many other languages (using native Unsigned Byte):
						//         0x20 / 2 = 0x10             //   32 / 2 =  16
						//         0x88 / 2 = 0x44             //  136 / 2 =  68
						//     In Java (without using the trick of the binary and):
						//         0x20 / 2 = 0x10             //   32 / 2 =  16
						//         0x88 / 2 = 0xc4             // -120 / 2 = -60  <=  Unexpected result
						//     In Java (using the trick of the binary and):
						//         (0x20 & 0xFF) / 2 = 0x10    //   32 / 2 =  16
						//         (0x88 & 0xFF) / 2 = 0x44    //  136 / 2 =  68
						reds[i] = (byte)((reds[i] & 0xFF) * ratioRed);
						greens[i] = (byte)((greens[i] & 0xFF) * ratioGreen);
						blues[i] = (byte)((blues[i] & 0xFF) * ratioBlue);
					}
				} else {
					// For robustest
					// There is a very unlikely case where each colour parts has different array size (corrupted image?).
					// This block do the same as the common scenario, but is 3x slower.
					for (int i=0; i<reds.length; i++) {
						reds[i] = (byte)((reds[i] & 0xFF) * ratioRed);
					}
					for (int i=0; i<greens.length; i++) {
						greens[i] = (byte)((greens[i] & 0xFF) * ratioGreen);
					}
					for (int i=0; i<blues.length; i++) {
						blues[i] = (byte)((blues[i] & 0xFF) * ratioBlue);
					}
				}

				// Create a new indexed image, using the new index colour model.
				IndexColorModel newModel;
				if (transparentPixel == -1) {
					newModel = new IndexColorModel(8, size, reds, greens, blues);
				} else {
					newModel = new IndexColorModel(8, size, reds, greens, blues, transparentPixel);
				}

				image = new BufferedImage(width, height, image.getType(), newModel);
				image.setData(raster);
			} else {
				// Not indexed images (every pixels has info about its colour components);
				//     PNG, JPEG, BMP, etc.

				WritableRaster raster = image.getRaster();
				// NOTE: Some image type (such as GIF) has only one colour (pixel) instead of 3 (for red, green and blue)
				//     or maybe 4 (including alpha). Those image types are not supported by this method so they are converted
				//     a supported format. If an untested image format is considered has been supported but in fact it is not,
				//     this check will prevent unnecessary parsing by returning the original image.
				if (raster.getPixel(0, 0, (int[]) null).length < 3) {
					LOGGER.log(Level.SEVERE, "Unsupported image type: {0}   Source: {1}", new Object[]{imageType, source});
					return image;
				}

				// Compute some values here instead of inside the loop, to save valuable CPU time
				float ratioRed = red / (float)0xFF,
					ratioGreen = green / (float)0xFF,
					ratioBlue = blue / (float)0xFF;

				for (int x = 0; x < width; x++) {
					for (int y = 0; y < height; y++) {
						int[] pixels = raster.getPixel(x, y, (int[]) null);
						if (pixels.length >= 3) {
							// The best bit-wise equivalent I could find quantise the values (pixels[0] &= red)
							pixels[0] *= ratioRed;
							pixels[1] *= ratioGreen;
							pixels[2] *= ratioBlue;
							raster.setPixel(x, y, pixels);
						}
					}
				}
			}
		} catch (Exception ex) {
			LOGGER.log(Level.SEVERE, "Could not convert the image: " + source, ex);
		}

		return image;
	}

	public static BufferedImage convert(BufferedImage src, int bufImgType) {
		BufferedImage img = new BufferedImage(src.getWidth(), src.getHeight(), bufImgType);
		Graphics2D g2d= img.createGraphics();
		g2d.drawImage(src, 0, 0, null);
		g2d.dispose();
		return img;
	}

	/**
	 * -----------------------------------------------------------------------------
	 * | Constant            | Value | File type                 | Supported       |
	 * -----------------------------------------------------------------------------
	 * | TYPE_CUSTOM         | 0     | Grayscale PNG             | Yes (converted) |
	 * | TYPE_INT_RGB        | 1     | BMP 32 bits (A8 R8 G8 B8) | Yes             |
	 * | TYPE_INT_ARGB       | 2     | ?                         | Probably        |
	 * | TYPE_INT_ARGB_PRE   | 3     | ?                         | Probably        |
	 * | TYPE_INT_BGR        | 4     | ?                         | Probably        |
	 * | TYPE_3BYTE_BGR      | 5     | JPEG                      | Yes             |
	 * | TYPE_4BYTE_ABGR     | 6     | PNG                       | Yes             |
	 * | TYPE_4BYTE_ABGR_PRE | 7     | ?                         | Probably        |
	 * | TYPE_USHORT_565_RGB | 8     | ? (*)                     | Yes             |
	 * | TYPE_USHORT_555_RGB | 9     | BMP 16 bits (X1 R5 G5 B5) | Yes             |
	 * | TYPE_BYTE_GRAY      | 10    | Grayscale JPEG            | Yes (converted) |
	 * | TYPE_USHORT_GRAY    | 11    | ?                         | Yes (converted) |
	 * | TYPE_BYTE_BINARY    | 12    | ?                         | Yes (converted) |
	 * | TYPE_BYTE_INDEXED   | 13    | GIF, Indexed PNG          | Yes             |
	 * -----------------------------------------------------------------------------
	 *
	 * (*) I though this one would be "BMP 16 bits (R5 G6 B5)", but java do not support this BMP format.
	 */
	private static boolean isSupported(int imageType) {
		return (imageType == BufferedImage.TYPE_INT_RGB ||        // Verified
				imageType == BufferedImage.TYPE_INT_ARGB ||
				imageType == BufferedImage.TYPE_INT_ARGB_PRE ||
				imageType == BufferedImage.TYPE_INT_BGR ||
				imageType == BufferedImage.TYPE_3BYTE_BGR ||      // Verified
				imageType == BufferedImage.TYPE_4BYTE_ABGR ||     // Verified
				imageType == BufferedImage.TYPE_4BYTE_ABGR_PRE ||
				imageType == BufferedImage.TYPE_USHORT_565_RGB || // Verified
				imageType == BufferedImage.TYPE_USHORT_555_RGB || // Verified
				imageType == BufferedImage.TYPE_BYTE_INDEXED);
	}

	private static int normaliseColour(int colour) {
		if (colour < 0) { return 0; }
		if (colour > 255) { return 255; }
		return colour;
	}
}
